#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：T1P.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/8 21:39
深度强化学习可以出甘特图比较好的方法-注释版本
'''
# !/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
FJSP with Deep Q-Network (DQN) — 解决20工件×4工序/工件×5机器的柔性作业车间调度问题
核心思路：通过DQN学习"状态→最优动作"的映射，最小化总完工时间（Makespan）
使用说明：
  1. 安装依赖：pip install torch numpy matplotlib
  2. 运行代码：python fjsp_dqn_20x4x5.py
"""

import random
from dataclasses import dataclass  # 用于定义结构化数据类，简化数据存储和访问
from typing import Dict, List, Tuple  # 类型提示，提高代码可读性
import numpy as np  # 数值计算库，用于状态表示、数组操作
import torch  # PyTorch深度学习框架，构建DQN网络
import torch.nn as nn  # 神经网络层模块
import torch.optim as optim  # 优化器模块
from collections import deque  # 双端队列，用于实现经验回放缓冲区
import matplotlib.pyplot as plt  # 绘图库，用于绘制甘特图和训练曲线


# ------------------------------
# 1. 问题定义模块（Problem Definition）
# 作用：定义FJSP问题的核心数据结构，生成随机调度实例
# ------------------------------

@dataclass  # 数据类，自动生成构造函数、__repr__等方法
class Operation:
    """工序（Operation）类：描述单个工序的可选机器和对应加工时间"""
    # 键：机器ID（整数），值：该机器加工此工序的时间（整数）
    machine_time: Dict[int, int]


@dataclass
class Job:
    """工件（Job）类：描述单个工件包含的所有工序"""
    # 工件的工序列表，按加工顺序排列（必须先完成前序工序才能执行后续工序）
    ops: List[Operation]


@dataclass
class Problem:
    """FJSP问题实例类：包含所有工件、机器数量，提供问题相关的统计信息"""
    # 所有工件的列表
    jobs: List[Job]
    # 机器总数
    m: int

    @property  # 装饰器，将方法转为属性，无需括号即可调用
    def n_jobs(self) -> int:
        """获取工件总数"""
        return len(self.jobs)

    @property
    def n_ops_total(self) -> int:
        """获取所有工件的总工序数（调度任务的总工作量）"""
        return sum(len(j.ops) for j in self.jobs)

    def max_ops_per_job(self) -> int:
        """获取单个工件的最大工序数（用于动作维度计算和状态归一化）"""
        return max(len(j.ops) for j in self.jobs)


def build_demo_problem(n_jobs: int, m: int, ops_per_job: int, seed: int = 7) -> Problem:
    """
    生成随机FJSP测试实例
    参数：
        n_jobs: 工件数量
        m: 机器数量
        ops_per_job: 每个工件的固定工序数
        seed: 随机种子（保证实例可复现）
    返回：
        Problem: FJSP问题实例
    生成规则：
        - 每个工序随机选择1-2台机器（避免可选机器过多导致动作空间过大）
        - 加工时间随机生成1-9（小范围整数，简化计算）
    """
    # 初始化随机数生成器，固定种子确保实例可复现
    rng = random.Random(seed)
    jobs: List[Job] = []  # 存储所有工件

    # 循环生成每个工件
    for _ in range(n_jobs):
        ops = []  # 存储当前工件的所有工序
        # 循环生成当前工件的每个工序
        for _o in range(ops_per_job):
            # 随机选择1-2台机器（若机器总数<2则选全部机器）
            machines = sorted(rng.sample(range(m), k=min(2, m)))
            # 为每台可选机器分配1-9的随机加工时间
            machine_time = {mm: rng.randint(1, 9) for mm in machines}
            # 将生成的工序加入工件的工序列表
            ops.append(Operation(machine_time))
        # 将生成的工件加入总工件列表
        jobs.append(Job(ops))

    # 返回完整的FJSP问题实例
    return Problem(jobs=jobs, m=m)


# ------------------------------
# 2. 环境模块（Environment）
# 作用：模拟FJSP调度过程，实现状态转换、动作执行、奖励计算，是智能体与问题交互的接口
# ------------------------------

class FJSPEnv:
    """
    FJSP环境类：
    - 维护调度系统的实时状态（机器就绪时间、工件就绪时间、已完成工序数等）
    - 提供动作编码/解码（将整数动作映射为（工件-工序-机器）三元组）
    - 执行动作并返回新状态、奖励、终止信号（是否完成所有调度）
    """

    def __init__(self, problem: Problem):
        """
        环境初始化
        参数：
            problem: FJSP问题实例（定义了工件、工序、机器的基础信息）
        """
        self.pb = problem  # 存储FJSP问题实例
        self.n_jobs = self.pb.n_jobs  # 工件总数
        self.m = self.pb.m  # 机器总数
        self.max_ops = self.pb.max_ops_per_job()  # 单个工件的最大工序数
        # 动作维度：每个动作对应（工件j, 工序k, 机器m）的组合，共n_jobs×max_ops×m个可能动作
        self.action_dim = self.n_jobs * self.max_ops * self.m
        self.reset()  # 初始化环境状态

    def reset(self) -> np.ndarray:
        """
        重置环境到初始状态（开始新的调度任务）
        返回：
            np.ndarray: 初始状态向量
        状态变量说明：
            next_op: 每个工件的"下一道待执行工序"索引（初始为0，即首道工序）
            job_ready: 每个工件的就绪时间（初始为0，即工件可立即开始首道工序）
            machine_ready: 每台机器的就绪时间（初始为0，即机器可立即加工）
            makespan: 当前总完工时间（初始为0，所有工序完成后为最终完工时间）
            done_ops: 已完成的工序总数（初始为0，达到总工序数时任务结束）
            assign: 已调度工序的记录（键：(工件ID, 工序ID)，值：(机器ID, 开始时间, 结束时间)）
        """
        self.next_op = np.zeros(self.n_jobs, dtype=np.int32)  # 工件j的下一道工序索引
        self.job_ready = np.zeros(self.n_jobs, dtype=np.float32)  # 工件j的就绪时间
        self.machine_ready = np.zeros(self.m, dtype=np.float32)  # 机器m的就绪时间
        self.makespan = 0.0  # 总完工时间（所有机器中最晚的结束时间）
        self.done_ops = 0  # 已完成的工序数
        self.assign = {}  # 工序调度记录
        return self._state()  # 返回初始状态

    def _state(self) -> np.ndarray:
        """
        构建当前环境的状态向量（用于输入DQN网络）
        状态向量组成（归一化处理，避免数值范围差异影响训练稳定性）：
            1. 机器就绪时间（m维）：每台机器的当前就绪时间 / 最大就绪时间（防止数值过大）
            2. 工件就绪时间（n_jobs维）：每个工件的当前就绪时间 / 最大就绪时间
            3. 下一道工序索引（n_jobs维）：每个工件的下一道工序索引 / 最大工序数（归一化到0-1）
        返回：
            np.ndarray: 归一化后的状态向量（长度 = m + n_jobs + n_jobs）
        """
        # 计算机器就绪时间的最大值（避免除以0，最小取1.0）
        mr_max = float(np.max(self.machine_ready)) if self.machine_ready.size > 0 else 1.0
        # 计算工件就绪时间的最大值
        jr_max = float(np.max(self.job_ready)) if self.job_ready.size > 0 else 1.0

        # 1. 机器就绪时间归一化（0-1范围）
        mr = self.machine_ready / (1.0 + max(mr_max, 1.0))
        # 2. 工件就绪时间归一化（0-1范围）
        jr = self.job_ready / (1.0 + max(jr_max, 1.0))
        # 3. 下一道工序索引归一化（0-1范围）
        noi = self.next_op.astype(np.float32) / max(float(self.max_ops), 1.0)

        # 拼接三个部分，形成最终状态向量
        return np.concatenate([mr, jr, noi], axis=0).astype(np.float32)

    def _decode_action(self, a: int) -> Tuple[int, int, int]:
        """
        将整数动作解码为（工件ID, 工序ID, 机器ID）三元组
        动作编码规则：a = j × (max_ops × m) + k × m + mm
        参数：
            a: 整数动作（0 ~ action_dim-1）
        返回：
            Tuple[int, int, int]: (工件j, 工序k, 机器mm)
        """
        # 每个工件对应的动作块大小（max_ops×m个动作）
        jm_block = self.max_ops * self.m
        j = a // jm_block  # 工件ID（整除块大小）
        rem = a % jm_block  # 剩余部分（块内偏移）
        k = rem // self.m  # 工序ID（块内整除机器数）
        mm = rem % self.m  # 机器ID（块内取余机器数）
        return int(j), int(k), int(mm)

    def _encode_action(self, j: int, k: int, mm: int) -> int:
        """
        将（工件ID, 工序ID, 机器ID）三元组编码为整数动作（_decode_action的逆操作）
        参数：
            j: 工件ID
            k: 工序ID
            mm: 机器ID
        返回：
            int: 整数动作
        """
        return j * (self.max_ops * self.m) + k * self.m + mm

    def valid_action(self, j: int, k: int, mm: int) -> bool:
        """
        判断（工件j, 工序k, 机器mm）是否为有效动作
        有效动作的4个条件：
            1. 工件ID j 在合法范围内（0 ≤ j < n_jobs）
            2. 工序k是工件j的下一道待执行工序（k == next_op[j]）
            3. 工序k在工件j的工序列表范围内（k < 工件j的总工序数）
            4. 机器mm是工序k的可选机器（mm在工序k的machine_time键中）
        参数：
            j: 工件ID
            k: 工序ID
            mm: 机器ID
        返回：
            bool: True（有效）/ False（无效）
        """
        if j < 0 or j >= self.n_jobs:  # 条件1：工件ID合法
            return False
        if k != int(self.next_op[j]):  # 条件2：工序是当前待执行的
            return False
        if k >= len(self.pb.jobs[j].ops):  # 条件3：工序ID在工件的工序范围内
            return False
        if mm not in self.pb.jobs[j].ops[k].machine_time:  # 条件4：机器是工序的可选机器
            return False
        return True

    def available_actions_mask(self) -> np.ndarray:
        """
        生成有效动作掩码（用于排除无效动作，避免智能体选择非法动作）
        返回：
            np.ndarray: 掩码向量（长度=action_dim），1表示有效动作，0表示无效动作
        """
        mask = np.zeros(self.action_dim, dtype=np.float32)
        # 遍历所有工件，找到每个工件的待执行工序及其可选机器
        for j in range(self.n_jobs):
            k = int(self.next_op[j])  # 工件j的待执行工序
            if k >= len(self.pb.jobs[j].ops):  # 工件j已完成所有工序，无有效动作
                continue
            # 遍历工序k的所有可选机器，标记对应动作为有效
            for mm in self.pb.jobs[j].ops[k].machine_time.keys():
                mask[self._encode_action(j, k, mm)] = 1.0
        return mask

    def step(self, action: int) -> Tuple[np.ndarray, float, bool, Dict]:
        """
        执行动作（调度一道工序），更新环境状态，返回交互结果
        参数：
            action: 整数动作（待执行的调度动作）
        返回：
            Tuple: (新状态向量, 奖励, 是否终止, 额外信息)
        核心逻辑：
            1. 解码动作，判断是否有效
            2. 有效动作：计算工序的开始/结束时间，更新机器、工件的就绪时间
            3. 计算奖励（惩罚完工时间增量，鼓励缩短总完工时间）
            4. 判断是否终止（所有工序完成）
        """
        # 1. 解码动作，获取（工件j, 工序k, 机器mm）
        j, k, mm = self._decode_action(action)

        # 2. 处理无效动作（给予惩罚，不更新状态）
        if not self.valid_action(j, k, mm):
            reward = -1.0  # 无效动作惩罚（避免智能体学习选无效动作）
            done = (self.done_ops == self.pb.n_ops_total)  # 检查是否意外终止（通常不会）
            return self._state(), reward, done, {}

        # 3. 处理有效动作：计算工序的开始和结束时间
        # 3.1 获取工序k在机器mm上的加工时间
        pt = float(self.pb.jobs[j].ops[k].machine_time[mm])
        # 3.2 工序开始时间 = max(机器mm的就绪时间, 工件j的就绪时间)（需等机器和工件都就绪）
        start = max(float(self.machine_ready[mm]), float(self.job_ready[j]))
        # 3.3 工序结束时间 = 开始时间 + 加工时间
        end = start + pt

        # 4. 更新环境状态
        self.machine_ready[mm] = end  # 更新机器mm的就绪时间为工序结束时间
        self.job_ready[j] = end  # 更新工件j的就绪时间为工序结束时间（下一道工序需等当前完成）
        self.assign[(j, k)] = (mm, start, end)  # 记录该工序的调度结果

        # 5. 更新总完工时间和已完成工序数
        old_makespan = self.makespan  # 记录更新前的完工时间
        self.makespan = max(self.makespan, end)  # 总完工时间取当前最大结束时间
        self.next_op[j] += 1  # 工件j的下一道待执行工序索引+1
        self.done_ops += 1  # 已完成工序数+1

        # 6. 计算奖励（关键设计：鼓励减少完工时间）
        # 奖励 = -（完工时间增量） - 0.01（步长惩罚，避免冗余步骤）
        # 完工时间增量越小，奖励越高；若未增加（end ≤ old_makespan），则增量惩罚为0
        reward = -(self.makespan - old_makespan) - 0.01

        # 7. 判断是否终止（所有工序都已完成）
        done = (self.done_ops == self.pb.n_ops_total)

        # 返回新状态、奖励、终止信号、空字典（额外信息）
        return self._state(), reward, done, {}


# ------------------------------
# 3. DQN智能体模块（DQN Agent）
# 作用：实现Deep Q-Network算法，包括Q网络、目标网络、经验回放、epsilon贪心策略
# ------------------------------

class QNet(nn.Module):
    """
    Q网络（价值网络）：输入状态，输出每个动作的Q值（动作价值）
    Q值含义：Q(s,a) = 状态s下选择动作a的长期累积奖励期望
    """

    def __init__(self, state_dim: int, action_dim: int):
        """
        初始化Q网络结构
        参数：
            state_dim: 状态向量维度（输入维度）
            action_dim: 动作维度（输出维度，每个动作对应一个Q值）
        网络结构：3层全连接网络（输入层→隐藏层1→隐藏层2→输出层）
            - 隐藏层：256个神经元，ReLU激活函数（引入非线性，拟合复杂映射）
            - 输出层：无激活函数（Q值可正可负，无需约束范围）
        """
        super().__init__()  # 调用父类nn.Module的构造函数
        self.net = nn.Sequential(  # 顺序执行的网络模块
            nn.Linear(state_dim, 256),  # 输入层→隐藏层1：state_dim→256
            nn.ReLU(),  # 隐藏层1激活函数：ReLU
            nn.Linear(256, 256),  # 隐藏层1→隐藏层2：256→256
            nn.ReLU(),  # 隐藏层2激活函数：ReLU
            nn.Linear(256, action_dim)  # 隐藏层2→输出层：256→action_dim（每个动作的Q值）
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播（计算Q值）
        参数：
            x: 输入状态张量（shape: [batch_size, state_dim]）
        返回：
            torch.Tensor: 每个动作的Q值（shape: [batch_size, action_dim]）
        """
        return self.net(x)


class DQNAgent:
    """
    DQN智能体类：实现DQN核心逻辑，包括动作选择、经验回放、网络训练、目标网络更新
    核心机制：
        - 经验回放（Experience Replay）：打破样本相关性，稳定训练
        - 目标网络（Target Network）：减少目标Q值波动，避免训练震荡
        - epsilon贪心策略（Epsilon-Greedy）：平衡探索（选随机动作）和利用（选Q值最大动作）
    """

    def __init__(
            self,
            state_dim: int,
            action_dim: int,
            lr: float = 1e-3,  # 学习率（控制参数更新步长）
            gamma: float = 0.99,  # 折扣因子（未来奖励的权重，0→只看即时奖励，1→同等看待未来奖励）
            eps_start: float = 1.0,  # epsilon初始值（全探索）
            eps_end: float = 0.05,  # epsilon最小值（最小探索率）
            eps_decay: float = 0.997,  # epsilon衰减率（每次训练后减少探索率）
            buffer_size: int = 150_000,  # 经验回放缓冲区大小（存储历史经验的最大数量）
            batch_size: int = 256,  # 批次大小（每次训练采样的经验数量）
            target_tau: float = 1.0,  # 目标网络更新参数（1.0→硬更新，<1.0→软更新）
    ):
        """初始化DQN智能体的参数、网络、优化器、经验缓冲区"""
        # 设备选择：优先使用GPU（CUDA），无GPU则用CPU
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # 1. 初始化Q网络（当前网络）和目标网络
        self.q = QNet(state_dim, action_dim).to(self.device)  # 当前网络：用于选动作和计算当前Q值
        self.q_target = QNet(state_dim, action_dim).to(self.device)  # 目标网络：用于计算目标Q值
        self.q_target.load_state_dict(self.q.state_dict())  # 初始时目标网络参数=当前网络参数

        # 2. 初始化优化器（Adam：常用的自适应学习率优化器）
        self.opt = optim.Adam(self.q.parameters(), lr=lr)

        # 3. 初始化DQN核心参数
        self.gamma = gamma  # 折扣因子
        self.eps = eps_start  # 当前epsilon值（探索率）
        self.eps_end = eps_end  # epsilon最小值
        self.eps_decay = eps_decay  # epsilon衰减率
        self.batch_size = batch_size  # 批次大小
        self.target_tau = target_tau  # 目标网络更新参数

        # 4. 初始化经验回放缓冲区（双端队列，超出大小自动删除旧经验）
        self.memory = deque(maxlen=buffer_size)

    def act(self, state: np.ndarray, valid_mask: np.ndarray) -> int:
        """
        选择动作（epsilon贪心策略）
        参数：
            state: 当前状态向量（用于计算Q值）
            valid_mask: 有效动作掩码（1=有效，0=无效）
        返回：
            int: 选择的整数动作
        策略逻辑：
            - 以概率eps：随机选择有效动作（探索未知动作）
            - 以概率1-eps：选择Q值最大的有效动作（利用已知最优动作）
        """
        # 探索：随机选择有效动作
        if np.random.rand() < self.eps:
            valid_indices = np.nonzero(valid_mask)[0]  # 找到所有有效动作的索引
            if len(valid_indices) == 0:  # 极端情况：无有效动作（理论上不会发生）
                return np.random.randint(valid_mask.size)
            return int(np.random.choice(valid_indices))  # 随机选一个有效动作

        # 利用：选择Q值最大的有效动作
        with torch.no_grad():  # 禁用梯度计算（仅前向传播，节省计算资源）
            # 将状态向量转为张量，添加batch维度（网络输入需为[1, state_dim]）
            s = torch.tensor(state, dtype=torch.float32, device=self.device).unsqueeze(0)
            # 计算所有动作的Q值，转为numpy数组（方便后续处理）
            qvals = self.q(s).squeeze(0).cpu().numpy()
            # 将无效动作的Q值设为极小值（-1e9），避免被选中
            qvals[valid_mask < 0.5] = -1e9
            # 返回Q值最大的动作索引
            return int(np.argmax(qvals))

    def remember(self, s: np.ndarray, a: int, r: float, s2: np.ndarray, d: bool, mask2: np.ndarray):
        """
        存储经验到回放缓冲区
        经验格式：(当前状态s, 动作a, 奖励r, 下一状态s2, 是否终止d, 下一状态的有效动作掩码mask2)
        参数：
            s: 当前状态
            a: 执行的动作
            r: 获得的奖励
            s2: 执行动作后的下一状态
            d: 是否终止（True=任务结束，False=继续）
            mask2: 下一状态的有效动作掩码
        """
        self.memory.append((s, a, r, s2, d, mask2))

    def train_step(self) -> float:
        """
        执行一次训练步骤（从缓冲区采样经验，更新Q网络和目标网络）
        返回：
            float: 本次训练的损失值（MSE损失）
        训练逻辑（Bellman方程）：
            目标Q值 = 即时奖励r + gamma × max_a'(Q_target(s2, a')) × (1 - d)
            （d=1时任务终止，未来奖励为0，故乘以(1-d)）
            损失 = MSE(当前Q值Q(s,a), 目标Q值)
        """
        # 缓冲区经验数量不足批次大小，无法训练（返回损失0）
        if len(self.memory) < self.batch_size:
            return 0.0

        # 1. 从缓冲区随机采样一批经验（打破样本时间相关性）
        batch = random.sample(self.memory, self.batch_size)

        # 2. 将经验转换为PyTorch张量（便于GPU计算）
        # 批量处理：将每个经验的对应字段堆叠成数组，再转为张量
        s = torch.tensor(np.stack([b[0] for b in batch]), dtype=torch.float32, device=self.device)
        a = torch.tensor([b[1] for b in batch], dtype=torch.int64, device=self.device).unsqueeze(1)
        r = torch.tensor([b[2] for b in batch], dtype=torch.float32, device=self.device).unsqueeze(1)
        s2 = torch.tensor(np.stack([b[3] for b in batch]), dtype=torch.float32, device=self.device)
        d = torch.tensor([b[4] for b in batch], dtype=torch.float32, device=self.device).unsqueeze(1)
        mask2 = torch.tensor(np.stack([b[5] for b in batch]), dtype=torch.float32, device=self.device)

        # 3. 计算当前Q值（Q(s,a)：当前网络对采样动作的Q值估计）
        # gather(1, a)：按动作a的索引，提取每个样本对应动作的Q值（shape: [batch_size, 1]）
        q_curr = self.q(s).gather(1, a)

        # 4. 计算目标Q值（基于目标网络，避免当前网络参数波动影响目标）
        with torch.no_grad():  # 目标网络不计算梯度（仅用于评估）
            # 4.1 计算下一状态s2的所有动作Q值（目标网络输出）
            q_next_all = self.q_target(s2)
            # 4.2 创建极小值张量（用于掩盖无效动作的Q值）
            very_neg = torch.full_like(q_next_all, -1e9)
            # 4.3 掩盖无效动作：有效动作保留Q值，无效动作设为-1e9
            q_next_all = torch.where(mask2 > 0.5, q_next_all, very_neg)
            # 4.4 取下一状态的最大Q值（max_a' Q_target(s2, a')）
            q_next_max = q_next_all.max(dim=1, keepdim=True).values
            # 4.5 按Bellman方程计算目标Q值
            target = r + (1.0 - d) * self.gamma * q_next_max

        # 5. 计算损失（均方误差MSE：当前Q值与目标Q值的差距）
        loss = nn.MSELoss()(q_curr, target)

        # 6. 更新当前Q网络（反向传播+梯度下降）
        self.opt.zero_grad()  # 清空之前的梯度（避免累积）
        loss.backward()  # 反向传播计算梯度
        nn.utils.clip_grad_norm_(self.q.parameters(), 1.0)  # 梯度裁剪（防止梯度爆炸）
        self.opt.step()  # 优化器更新网络参数

        # 7. 更新目标网络（硬更新或软更新）
        if self.target_tau >= 1.0:
            # 硬更新：直接将当前网络参数复制到目标网络（适用于固定频率更新）
            self.q_target.load_state_dict(self.q.state_dict())
        else:
            # 软更新：目标网络参数 = 目标网络参数×(1-tau) + 当前网络参数×tau（平滑更新）
            for tp, p in zip(self.q_target.parameters(), self.q.parameters()):
                tp.data.copy_(tp.data * (1.0 - self.target_tau) + p.data * self.target_tau)

        # 8. 更新epsilon（衰减探索率，随着训练增加利用比例）
        self.eps = max(self.eps_end, self.eps * self.eps_decay)

        # 返回本次训练的损失值（转为Python浮点数）
        return float(loss.item())


# ------------------------------
# 4. 训练与评估模块（Training & Evaluation）
# 作用：实现智能体的训练循环、策略评估、调度结果可视化（甘特图）
# ------------------------------

def train_dqn(env: FJSPEnv, episodes: int = 1500) -> Tuple[DQNAgent, List[Tuple[int, float, float, float]]]:
    """
    训练DQN智能体
    参数：
        env: FJSP环境实例
        episodes: 训练轮次（每个轮次对应一次完整的调度任务）
    返回：
        Tuple: (训练好的智能体, 训练历史记录)
        历史记录格式：(轮次ep, 总奖励ep_reward, 总损失ep_loss, 完工时间makespan)
    """
    # 1. 获取状态维度，初始化智能体
    state_dim = env._state().size  # 状态向量的长度
    agent = DQNAgent(
        state_dim=state_dim,
        action_dim=env.action_dim,
        lr=1e-3,  # 学习率
        gamma=0.99,  # 折扣因子（重视未来奖励）
        eps_start=1.0,  # 初始全探索
        eps_end=0.05,  # 最小探索率5%
        eps_decay=0.997,  # 每轮epsilon衰减3%
        buffer_size=150000,  # 大缓冲区存储更多经验
        batch_size=256,  # 批次大小
        target_tau=1.0,  # 硬更新（每步训练后更新目标网络）
    )

    # 2. 训练历史记录（用于后续分析）
    history = []

    # 3. 循环训练每一轮次（episode）
    for ep in range(1, episodes + 1):
        # 3.1 重置环境，获取初始状态和有效动作掩码
        s = env.reset()
        mask = env.available_actions_mask()

        # 3.2 初始化本轮次的统计变量
        ep_reward = 0.0  # 本轮总奖励
        ep_loss = 0.0  # 本轮总损失
        steps = 0  # 本轮步数（调度的工序数）

        # 3.3 单轮次调度循环（直到所有工序完成）
        while True:
            # 选择动作
            a = agent.act(s, mask)
            # 执行动作，获取交互结果
            s2, r, done, _ = env.step(a)
            # 获取下一状态的有效动作掩码
            mask2 = env.available_actions_mask()
            # 存储经验到缓冲区
            agent.remember(s, a, r, s2, done, mask2)
            # 执行一次训练步骤，累积损失
            l = agent.train_step()
            ep_loss += l
            # 累积奖励
            ep_reward += r
            # 步数+1
            steps += 1
            # 更新当前状态和掩码（准备下一步）
            s, mask = s2, mask2

            # 若所有工序完成，结束本轮次
            if done:
                break

        # 3.4 记录本轮次的训练结果
        history.append((ep, ep_reward, ep_loss, env.makespan))

        # 3.5 定期打印训练进度（每50轮或第1轮）
        if ep % 50 == 0 or ep == 1:
            print(
                f"Episode {ep:4d} | steps={steps:4d} | reward={ep_reward:9.3f} | "
                f"loss={ep_loss:9.3f} | makespan={env.makespan:7.2f} | eps={agent.eps:5.3f}"
            )

    # 返回训练好的智能体和训练历史
    return agent, history


def evaluate_policy(env: FJSPEnv, agent: DQNAgent, greedy: bool = True) -> Tuple[Dict, float, np.ndarray]:
    """
    评估智能体的调度策略（无探索，纯利用）
    参数：
        env: FJSP环境实例
        agent: 训练好的DQN智能体
        greedy: 是否使用贪心策略（True=选Q值最大动作，False=保留探索）
    返回：
        Tuple: (调度记录assign, 总完工时间makespan, 机器就绪时间machine_ready)
    """
    # 保存当前epsilon（评估后恢复，避免影响后续训练）
    prev_eps = agent.eps
    # 贪心评估：将epsilon设为0（不探索，只选最优动作）
    if greedy:
        agent.eps = 0.0

    # 重置环境，获取初始状态和掩码
    s = env.reset()
    mask = env.available_actions_mask()

    # 执行调度循环（直到所有工序完成）
    while True:
        # 选择动作（贪心策略）
        a = agent.act(s, mask)
        # 执行动作
        s, r, done, _ = env.step(a)
        # 更新掩码
        mask = env.available_actions_mask()
        # 任务完成，退出循环
        if done:
            break

    # 恢复评估前的epsilon
    agent.eps = prev_eps

    # 返回调度结果：调度记录、完工时间、机器就绪时间
    return env.assign, env.makespan, env.machine_ready.copy()


def plot_gantt(pb: Problem, assign: Dict[Tuple[int, int], Tuple[int, float, float]], title: str = "Schedule"):
    """
    绘制甘特图（可视化调度结果）
    参数：
        pb: FJSP问题实例（用于获取机器数量）
        assign: 调度记录（键：(工件j, 工序k)，值：(机器mm, 开始时间s, 结束时间e)）
        title: 图表标题
    甘特图含义：
        - 纵轴：机器ID（M0, M1, ...）
        - 横轴：时间
        - 每个水平条：一道工序的加工时间（标注"Jj-Ok"表示工件j的工序k）
    """
    # 按机器分组存储工序（键：机器ID，值：[(开始时间, 持续时间, 工序标签)]）
    bars = {mm: [] for mm in range(pb.m)}
    for (j, k), (mm, s, e) in assign.items():
        # 计算持续时间 = 结束时间 - 开始时间
        duration = e - s
        # 工序标签："Jj-Ok"（如J0-O1表示工件0的工序1）
        label = f"J{j}-O{k}"
        bars[mm].append((s, duration, label))

    # 创建图表
    fig, ax = plt.subplots(figsize=(12, 5))  # 图表大小：12英寸宽，5英寸高

    # 遍历每个机器，绘制水平条
    yticks = []  # 纵轴刻度位置
    ylabels = []  # 纵轴刻度标签（机器ID）
    for idx, mm in enumerate(sorted(bars.keys())):
        # 按开始时间排序当前机器的工序（确保甘特图按时间顺序显示）
        items = sorted(bars[mm], key=lambda x: x[0])
        # 绘制每个工序的水平条
        for (start, dur, label) in items:
            # ax.barh(纵轴位置, 宽度(持续时间), 左侧起始位置(开始时间))
            ax.barh(idx, dur, left=start)
            # 在水平条中间添加工序标签（字体大小8，水平/垂直居中）
            ax.text(start + dur / 2, idx, label, ha="center", va="center", fontsize=8)
        # 记录纵轴刻度（每个机器对应一个刻度）
        yticks.append(idx)
        ylabels.append(f"M{mm}")

    # 设置图表属性
    ax.set_yticks(yticks)  # 纵轴刻度位置
    ax.set_yticklabels(ylabels)  # 纵轴刻度标签（机器ID）
    ax.set_xlabel("Time")  # 横轴标签（时间）
    ax.set_title(title)  # 图表标题
    plt.tight_layout()  # 自动调整布局，避免标签被截断
    plt.show()  # 显示图表


# ------------------------------
# 5. 主函数（Main）
# 作用：整合所有模块，生成问题实例、训练智能体、评估策略、绘制结果
# ------------------------------

def main():
    """主函数：执行完整的FJSP调度任务（生成实例→训练→评估→可视化）"""
    # 1. 生成FJSP问题实例（20工件，5机器，每个工件4道工序，种子7确保可复现）
    # pb = build_demo_problem(n_jobs=20, m=5, ops_per_job=4, seed=7)
    pb = build_demo_problem(n_jobs=50, m=10, ops_per_job=6, seed=7)
    # 2. 初始化FJSP环境
    env = FJSPEnv(pb)
    # 3. 训练DQN智能体（1500轮次）
    agent, _ = train_dqn(env, episodes=1500)

    # 4. 评估训练好的智能体（贪心策略，无探索）
    assign, makespan, _ = evaluate_policy(env, agent, greedy=True)

    # 5. 输出评估结果
    print(f"\nFinal greedy makespan: {makespan:.2f}")  # 打印最终完工时间
    # 打印每个工序的调度详情（按工件和工序排序）
    for (j, k), (mm, s, e) in sorted(assign.items()):
        print(f"Job {j} Op {k} -> M{mm}, start={s:.1f}, end={e:.1f}")

    # 6. 绘制甘特图（可视化调度结果）
    plot_gantt(
        pb,
        assign,
        title=f"FJSP DQN Schedule (20 Jobs × 4 Ops/Job × 5 Machines, Makespan={makespan:.1f})"
    )


# 程序入口：当脚本直接运行时，执行主函数
if __name__ == "__main__":
    main()